import idautils
# import idaapi
import idc

import re
import string
import pickle as pk
import idaapi
import logging
import sys
import os

sys.path.append("E:\\Anaconda3\\lib\\site-packages")
import pyparsing as pp  # noqa

FUNCATTR_START = 0  # function start address
FUNCATTR_END = 4  # function end address


def parse_one_line(line):
    """Parsing the oprands

    Args:
        line (str): The oprands string

    Returns:
        list: A token list
    """
    if "ymmword ptr" in line:
        line = line.replace("ymmword ptr", "")
    if "zmmword ptr" in line:
        line = line.replace("zmmword ptr", "")
    if "qword ptr" in line:
        line = line.replace("qword ptr", "")
    if "dword ptr" in line:
        line = line.replace("dword ptr", "")
    if "word ptr" in line:
        line = line.replace("word ptr", "")
    if "byte ptr" in line:
        line = line.replace("byte ptr", "")

    line = line.strip()
    tokens = []
    nums = pp.Word(pp.nums)
    minus = pp.oneOf("-")
    hexnums = minus * (0, 1) + (pp.Combine("0x" + pp.Word(pp.hexnums)) *
                                (0, 1) + pp.Combine(pp.Word(pp.hexnums) + "h") * (0, 1))
    s_nums = pp.Combine("k" + nums)
    numbers = hexnums * (0, 1) + nums * (0, 1)
    arithOp = pp.oneOf("+ - * / ,")
    left = pp.oneOf("[ {")
    right = pp.oneOf("] }")
    single_reg = pp.Combine(
        pp.Word(pp.alphas) + nums * (0, 2) + pp.Word(pp.alphas) * (0, 1) + ("." + numbers + pp.Word(pp.alphas)) *
        (0, 1))
    reg = single_reg*(0, 1) \
        + (left + single_reg + ((arithOp + single_reg)*(0, 2) + (arithOp + numbers)*(0, 1))*(0, 2) + right)*(0, 1) + ("{" + s_nums + "}")*(0, 1) + ("[" + numbers + "]")*(0, 1)
    # external_func = pp.Combine("_" + pp.Word(pp.alphas))
    # reg_or_numbers = reg * (0, 1) + numbers * (0, 1) + external_func * (0, 1)
    reg_or_numbers = reg * (0, 1) + numbers * (0, 1)
    expr = reg_or_numbers + ("," + reg_or_numbers) * (0, 3)
    tokens = [x for x in expr.parseString(line)]
    return tokens


def test(f):
    type_arr = []
    if idc.isCode(f):
        type_arr.append("code")

    if idc.isData(f):
        type_arr.append("data")

    if idc.isTail(f):
        type_arr.append("tail")

    if idc.isUnknown(f):
        type_arr.append("unknown")

    if idc.isHead(f):
        type_arr.append("head")

    print("Type is {}".format(type_arr))


jump_op = [
    "jz", "jnz", "jc", "jnc", "jo", "jno", "js", "jns", "jp", "jnp", "je", "jne", "jcxz", "jecxz", "jrcxz", "ja",
    "jnbe", "jb", "jae", "jbe", "jnae", "jnb", "jna", "jg", "jl", "jge", "jle", "jnle", "jnge", "jnl", "jng", "jmp"
]


def is_hex(s):
    hex_digits = set(string.hexdigits)
    return all(c in hex_digits for c in s)


def print_inline_function(addr):
    count = 0
    # from_addr = None
    # for i in idautils.XrefsTo(addr):
    #     print("From: {} To: {}".format(hex(i.frm), hex(i.to)))
    #     from_addr = i.to
    if addr is not None:
        for ea in idautils.FuncItems(addr):
            print(idc.GetDisasm(ea))
            count += 1
        print("Insert {} lines".format(count))
    else:
        print("Fail to inline function")


def get_tokens_one_line(line):
    tokens = []
    opcode = line.split(" ")[0]
    tokens.append(opcode)
    oprands = parse_one_line(line[len(opcode):].strip())

    # Handle some external function like _logf, _expf
    if opcode == "call" and line[len(opcode):].strip().startswith("_"):
        oprands.append(line[len(opcode):].strip()[1:])

    for opr in oprands:

        if opr.startswith("sub"):
            opr = "addr"

        if opr.isdigit() or "0x" in opr or (opr[-1] == 'h' and is_hex(opr[:-1])):
            if opr.isdigit():
                tokens.append(opr)
            elif "0x" in opr:
                if "00000000" in opr:
                    opr = opr.replace("00000000", "")

                if abs(int(opr, 16)) <= 2**15:
                    tokens.append(opr)
                else:
                    tokens.append("num")
            elif is_hex(opr[:-1]) and len(opr[:-1]) > 0:
                tmp = "0x{}".format(opr[:-1])
                if abs(int(tmp, 16)) <= 2**15:
                    tokens.append(tmp)
                else:
                    tokens.append("num")
        elif opr == ",":
            continue
        else:
            tokens.append(opr)
    return tokens


def get_function_blocks(function_items, funcname_addr_dict):

    inline_addrs = []
    inline_addrs_idx = []
    blocks = [[]]
    disasm_lines = []
    disasm_line_addrs = []
    splits_points = []  # The split points for spliting basic blocks
    block_edges = []
    addr_edges = []

    for idx, ea in enumerate(function_items):
        disasm_line = idc.GetDisasm(ea)
        opcode = disasm_line.split(" ")[0]
        tmp = idautils.DecodeInstruction(ea)
        if ";" in disasm_line:
            disasm_line = disasm_line.split(";")[0]
        if "#" in disasm_line:
            disasm_line = disasm_line.replace("#", "")
        if "@PAGEOFF" in disasm_line:
            disasm_line = disasm_line.replace("@PAGEOFF", "")
        if "@PAGE" in disasm_line:
            disasm_line = disasm_line.replace("@PAGE", "")

        if "var_s" in disasm_line:
            disasm_line = disasm_line.replace("var_s", "0x")

        if "var_" in disasm_line:
            disasmasm_line = disasm_line.replace("var_", "0x")

        if "cs:__" in disasm_line and "_ptr" in disasm_line:
            ptr_name = disasm_line.split("_")[-2]
            disasm_line = disasm_line.replace("cs:__{}_ptr".format(ptr_name), ptr_name)

        if "cs:" in disasm_line and "cs:__" not in disasm_line:
            p = re.compile(r'cs:[a-zA-Z]+')
            p_with_number = re.compile(r'cs:[a-zA-Z]+\_[0-9]+')
            name = p.search(disasm_line).group()
            disasm_line = re.sub(p_with_number, name.replace(":", ""), disasm_line)

        if "__TVM" in disasm_line and "_ptr" in disasm_line:
            disasm_line = disasm_line.replace("__", "")
            disasm_line = disasm_line.replace("_ptr", "")

        if ".logf" in disasm_line:
            disasm_line = disasm_line.replace(".logf", "logf")

        if ".expf" in disasm_line:
            disasm_line = disasm_line.replace(".expf", "expf")

        if opcode in jump_op:
            splits_points.append((idx, str(hex(tmp.Op1.addr))))
            # Change "jmp short loc_012" to "jmp 0x00012"
            disasm_line = "{} {}".format(opcode, str(hex(tmp.Op1.addr)))
            addr_edges.append([str(hex(ea)), str(hex(tmp.Op1.addr))])

        disasm_lines.append(disasm_line)
        disasm_line_addrs.append(str(hex(ea)))

        print(disasm_line)
        # X86 Pattern
        if "rdi," in disasm_line and "lea" in disasm_line:
            tmp = idautils.DecodeInstruction(ea)
            # print(get_operand_value(ea, 1))
            if tmp.Op2.addr != 0:
                inline_addrs.append(str(hex(tmp.Op2.addr)))
                inline_addrs_idx.append(idx)
        # ARM Pattern
        if "ADRP" in disasm_line and "sub_" in disasm_line:
            tmp = idautils.DecodeInstruction(ea)
            # print(get_operand_value(ea, 1))
            tmp_addr = disasm_line.split(",")[1].strip()
            if tmp_addr in funcname_addr_dict:
                inline_addrs.append(str(hex(funcname_addr_dict[tmp_addr])))
                inline_addrs_idx.append(idx)

    # Split basic blocks by jump operations
    new_splits_points = set()
    for sp in splits_points:
        pos, addr = sp
        new_splits_points.add(pos + 1)
        if addr in disasm_line_addrs:
            new_splits_points.add(disasm_line_addrs.index(addr))
        else:
            logging.error("Outer Addr: {}".format(addr))

    splits_points = new_splits_points

    block_boundary_start = []
    block_boundary_end = []
    for idx, line in enumerate(disasm_lines):
        if idx in splits_points:
            block_boundary_end.append(idx - 1)
            blocks.append([])
        if len(blocks[-1]) == 0:
            block_boundary_start.append(idx)
        blocks[-1].extend(get_tokens_one_line(line))
    block_boundary_end.append(len(disasm_lines) - 1)

    block_edges = []
    # retn check
    for i in range(len(block_boundary_start) - 1):
        tmp_edge = [i, i + 1]
        if "retn" in blocks[i] or "RET" in blocks[i]:
            continue
        block_edges.append(tmp_edge)

    for addr_edge in addr_edges:
        from_addr, to_addr = addr_edge
        if to_addr not in disasm_line_addrs or from_addr not in disasm_line_addrs:#判断跳转的地址在该函数里面
            continue
        from_bblock_idx = block_boundary_end.index(disasm_line_addrs.index(from_addr))
        to_bblock_idx = block_boundary_start.index(disasm_line_addrs.index(to_addr))
        block_edges.append([from_bblock_idx, to_bblock_idx])

    filtered_inline_addrs = []
    filtered_inline_idx = []
    for inline_addr, inline_idx in zip(inline_addrs, inline_addrs_idx):
        if inline_addr not in disasm_line_addrs:
            filtered_inline_addrs.append(inline_addr)
            filtered_inline_idx.append(inline_idx)

    inline_bblock_pos = []
    for inline_pos in filtered_inline_idx:
        for i in range(len(block_boundary_start)):
            st, ed = block_boundary_start[i], block_boundary_end[i]
            if inline_pos >= st and inline_pos <= ed:
                inline_bblock_pos.append(i)

    return blocks, filtered_inline_addrs, block_edges, inline_bblock_pos


# SAVE_DIR = "D:\\\\Research\\TVM_DATA\\results"
SAVE_DIR = "/Users/mhli/workspace/tvm_data"


def analysis():
    filename = idaapi.get_root_filename()
    data = {}
    funcname_addr_dict = {}
    for s_ea in idautils.Segments():
        for i, func in enumerate(idautils.Functions(idc.SegStart(s_ea), idc.SegEnd(s_ea))):
            func_name = idc.GetFunctionName(func)
            funcname_addr_dict[func_name] = func

    for s_ea in idautils.Segments():
        for i, func in enumerate(idautils.Functions(idc.SegStart(s_ea), idc.SegEnd(s_ea))):
            func_name = idc.GetFunctionName(func)
            logging.info("{0} function name is: {1}".format(i, func_name))
            print("{0} function name is: {1}".format(i, func_name))
            if "compute" in func_name:
                data[func_name] = {}
                function_items = idautils.FuncItems(func)
                data[func_name]["blocks"], inline_addrs, data[func_name][
                    "edges"], inline_bblock_pos = get_function_blocks(function_items, funcname_addr_dict)
                logging.info(inline_bblock_pos)
                logging.info(inline_addrs)
                logging.info(data[func_name]["edges"])
                data[func_name]["inline_addr_from_idx"] = inline_bblock_pos
                data[func_name]["inline_addr_bb_range"] = []
                for inline_addr in inline_addrs:
                    tmp_blocks, tmp_inline_addrs, tmp_edges, _ = get_function_blocks(
                        idautils.FuncItems(int(inline_addr[:-1], 16)), funcname_addr_dict)
                    for i in range(len(tmp_edges)):
                        tmp_edges[i][0] += len(data[func_name]["blocks"])
                        tmp_edges[i][1] += len(data[func_name]["blocks"])
                    data[func_name]["edges"].extend(tmp_edges)
                    data[func_name]["blocks"].extend(tmp_blocks)
                    inline_addr_range = (len(data[func_name]["blocks"]) - len(tmp_blocks) - 1,
                                         len(data[func_name]["blocks"]) - 1)
                    data[func_name]["inline_addr_bb_range"].append(inline_addr_range)

                    if len(tmp_inline_addrs) != 0:
                        logging.error("inline function has inline addr {}!".format(tmp_inline_addrs))
            # print("Op Type: {}".format(idc.GetOpType(ea, 1)))
    if not os.path.exists(SAVE_DIR):
        os.makedirs(SAVE_DIR)
    with open("{}\\{}.pkl".format(SAVE_DIR, filename), "wb") as f:
        pk.dump(data, f)

    pass


def main():
    """
    """
    idc.Wait()
    analysis()
    idc.Exit(0)


if __name__ == "__main__":
    main()
